import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.optim.lr_scheduler import StepLR
from network import resnet50

# import torchvision.models as models

import os
import time
import shutil
import random
import argparse
import numpy as np
from enum import Enum
from loguru import logger
from datetime import datetime
import matplotlib.pyplot as plt

logger.add(f"../../log/{str(datetime.today())[:10]}_ResNet_train.log", level='INFO', encoding="utf-8", enqueue=True)


def parse_option():
    parser = argparse.ArgumentParser('ResNet')
    parser.add_argument('--dummy', type=bool, help='Whether data needs to be generated', default=False)
    parser.add_argument('--epochs', type=int, help='the number of epoch', default=10)
    parser.add_argument('--batch_size', type=int, help='batch_size', default=32)
    parser.add_argument("--workers", type=int, help="the number of workers", default=8)
    parser.add_argument("--data_path", type=str, help="the path of model saved",
                        default='../../datasets/flowers')
    parser.add_argument("--save_path", type=str, help="the path of model saved",
                        default='../../models/ResNet.pth')
    parser.add_argument('--lr', default=0.0002, type=float, metavar='LR',
                        help='initial learning rate', dest='lr')
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')
    parser.add_argument('--print_freq', default=30, type=int,
                        metavar='W', help='Print frequency')
    args = parser.parse_args()  # 也可直接使用 args, _ = parser.parse_known_args()
    return args


def main(args):
    args.device = torch.device("cuda:7" if torch.cuda.is_available() else "cpu")
    logger.info("using {} device.".format(args.device))

    # Data loading code
    if args.dummy:
        print("=> Dummy data is used!")
        train_dataset = datasets.FakeData(1281167, (3, 224, 224), 1000, transforms.ToTensor())
        val_dataset = datasets.FakeData(50000, (3, 224, 224), 1000, transforms.ToTensor())
    else:
        traindir = os.path.join(args.data_path, 'train')
        valdir = os.path.join(args.data_path, 'test')
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        train_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))

        val_dataset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ]))

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    # save:  daisy,dandelion,roses,sunflowers,tulips
    classes_list = train_dataset.classes
    with open('../../datasets/flowers/classes.txt', 'w') as text_cls:
        classes_str = ','.join(classes_list)
        text_cls.write(classes_str)

    train_num = len(train_loader)  # 训练总数/batch_size
    val_num = len(val_loader)
    logger.info(f"using {train_num} batchs for training, {val_num} batchs for validation.")

    model = resnet50(num_classes=5)
    model.to(args.device)
    criterion = nn.CrossEntropyLoss().to(args.device)

    # # load pretrain weights
    # # download url: https://download.pytorch.org/models/resnet34-333f7ec4.pth
    # model_weight_path = "./resnet34-pre.pth"
    # assert os.path.exists(model_weight_path), "file {} does not exist.".format(model_weight_path)
    # model.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
    # # for param in net.parameters():
    # #     param.requires_grad = False
    # # change fc layer structure
    # in_channel = model.fc.in_features
    # model.fc = nn.Linear(in_channel, 5)  # change last layer for this project
    # model.to(args.device)

    optimizer = optim.Adam(model.parameters(), args.lr)

    # """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    # scheduler = StepLR(optimizer, step_size=50, gamma=0.1)

    best_acc1 = 0

    for epoch in range(args.epochs):  # loop over the dataset multiple times

        # train for one epoch
        train(train_loader, model, criterion, optimizer, epoch, args)

        # evaluate on validation set
        acc1 = validate(val_loader, model, criterion, args)

        # scheduler.step()

        # remember best acc@1 and save checkpoint
        is_best = acc1 > best_acc1
        best_acc1 = max(acc1, best_acc1)
        save_checkpoint({'epoch': epoch + 1,
                         'state_dict': model.state_dict(),
                         'best_acc1': best_acc1,
                         #    'optimizer' : optimizer.state_dict(),
                         #     'scheduler' : scheduler.state_dict()
                         },
                        is_best, args.save_path)

    logger.info('Finished Training')
    # torch.save(net.state_dict(), args.save_path)


def train(train_loader, model, criterion, optimizer, epoch, args):
    batch_time = AverageMeter('Time_train', ':6.3f')
    data_time = AverageMeter('Time_Date', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top3 = AverageMeter('Acc@3', ':6.2f')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses, top1, top3],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # move data to the same device as model
        images = images.to(args.device, non_blocking=True)
        target = target.to(args.device, non_blocking=True)

        # compute output
        output = model(images)
        loss = criterion(output, target)

        # measure accuracy and record loss
        acc1, acc3 = accuracy(output, target, topk=(1, 3))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top3.update(acc3[0], images.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:  # 每隔 print_freq 个batch_size打印一次训练结果
            progress.display(i + 1)


def validate(val_loader, model, criterion, args):
    batch_time = AverageMeter('Time_val', ':6.3f', Summary.NONE)
    losses = AverageMeter('Loss', ':.4e', Summary.NONE)
    top1 = AverageMeter('Acc@1', ':6.2f', Summary.AVERAGE)
    top3 = AverageMeter('Acc@3', ':6.2f', Summary.AVERAGE)
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, losses, top1, top3],
        prefix='Test: ')

    def run_validate(loader, base_progress=0):
        model.eval()
        with torch.no_grad():
            end = time.time()
            for i, (images, target) in enumerate(loader):
                i = base_progress + i

                # move data to the same device as model
                images = images.to(args.device, non_blocking=True)
                target = target.to(args.device, non_blocking=True)

                # compute output
                output = model(images)
                loss = criterion(output, target)

                # measure accuracy and record loss
                acc1, acc3 = accuracy(output, target, topk=(1, 3))
                losses.update(loss.item(), images.size(0))
                top1.update(acc1[0], images.size(0))
                top3.update(acc3[0], images.size(0))

                # measure elapsed time
                batch_time.update(time.time() - end)
                end = time.time()

    # switch to evaluate mode
    model.eval()

    run_validate(val_loader)

    progress.display_summary()

    return top1.avg


def save_checkpoint(state, is_best, filename='checkpoint.pth'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, '../../models/ResNet_best.pth')


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)
        sh = output.shape
        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


class Summary(Enum):
    NONE = 0
    AVERAGE = 1
    SUM = 2
    COUNT = 3


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self, name, fmt=':f', summary_type=Summary.AVERAGE):
        self.name = name
        self.fmt = fmt
        self.summary_type = summary_type
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

    def summary(self):
        fmtstr = ''
        if self.summary_type is Summary.NONE:
            fmtstr = ''
        elif self.summary_type is Summary.AVERAGE:
            fmtstr = '{name} {avg:.3f}'
        elif self.summary_type is Summary.SUM:
            fmtstr = '{name} {sum:.3f}'
        elif self.summary_type is Summary.COUNT:
            fmtstr = '{name} {count:.3f}'
        else:
            raise ValueError('invalid summary type %r' % self.summary_type)

        return fmtstr.format(**self.__dict__)


class ProgressMeter(object):

    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        print('\t'.join(entries))  # print(f"\t".join("ddd")) --> d	d	d

    def display_summary(self):
        entries = [" *** summary:"]
        entries += [meter.summary() for meter in self.meters]
        print(' '.join(entries))

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'


if __name__ == '__main__':
    args_ = parse_option()
    main(args_)
